package cn.felord.security.autoconfigure.jwt;

import cn.felord.security.autoconfigure.SecureUser;
import cn.felord.security.autoconfigure.TokenGenerator;
import cn.felord.security.autoconfigure.util.AuthenticationUtil;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.endpoint.OAuth2AccessTokenResponse;
import org.springframework.security.oauth2.jose.jws.SignatureAlgorithm;
import org.springframework.security.oauth2.jwt.JwsHeader;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.jwt.JwtEncoder;
import org.springframework.security.oauth2.jwt.JwtEncoderParameters;

import java.time.Clock;
import java.time.Instant;
import java.time.ZoneId;
import java.util.Collection;
import java.util.Collections;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * The type Jwt token generator.
 *
 * @author felord.cn
 * @since 2021 /8/7 16:48
 */
public class JwtTokenGenerator implements TokenGenerator<OAuth2AccessTokenResponse> {
    private final JwtEncoder jwtEncoder;
    private final JwtProperties jwtProperties;
    private JwtClaimsConsumer claimsConsumer = claimsMap -> {
    };

    /**
     * Instantiates a new Jwt token generator.
     *
     * @param jwtEncoder    the jwt encoder
     * @param jwtProperties the jwt properties
     */
    public JwtTokenGenerator(JwtEncoder jwtEncoder, JwtProperties jwtProperties) {
        this.jwtEncoder = jwtEncoder;
        this.jwtProperties = jwtProperties;
    }

    public JwtTokenGenerator(JwtEncoder jwtEncoder, JwtProperties jwtProperties, JwtClaimsConsumer claimsConsumer) {
        this.jwtEncoder = jwtEncoder;
        this.jwtProperties = jwtProperties;
        this.claimsConsumer = claimsConsumer;
    }

    @Override
    public OAuth2AccessTokenResponse tokenResponse(SecureUser secureUser) {
        JwsHeader jwsHeader = JwsHeader.with(SignatureAlgorithm.RS256).type("JWT").build();

        Instant issuedAt = Clock.system(ZoneId.of("Asia/Shanghai")).instant();
        Collection<GrantedAuthority> authorities = secureUser.getAuthorities();
        Set<String> scopes = authorities
                .stream()
                .map(GrantedAuthority::getAuthority)
                .collect(Collectors.toSet());
        long expiresAt = jwtProperties.getExpiresAt().getSeconds();

        String username = secureUser.getUsername();
        JwtClaimsSet jwtClaimsSet = JwtClaimsSet.builder()
                .issuer(jwtProperties.getIssuer())
                // 面向的用户群体
                .subject(username)
                // 接收jwt的一方 暂时为username
                .audience(Collections.singletonList(username))
                .claim("scope", scopes)
                .issuedAt(issuedAt)
                .expiresAt(issuedAt.plusSeconds(expiresAt))
                .claim(AuthenticationUtil.UID_CLAIM, secureUser.getUserId())
                .claim(AuthenticationUtil.CID_CLAIM, secureUser.getClientId())
                .claims(claimsConsumer::consumer)
                .build();

        Jwt accessToken = jwtEncoder.encode(JwtEncoderParameters.from(jwsHeader, jwtClaimsSet));
        // 移除 权限  避免 refresh token 用于请求接口
        Jwt refreshToken = jwtEncoder.encode(JwtEncoderParameters.from(jwsHeader, JwtClaimsSet.from(jwtClaimsSet)
                .claim("scope", Collections.emptySet())
                .expiresAt(issuedAt.plusSeconds(expiresAt * 2)).build()));

        return OAuth2AccessTokenResponse.withToken(accessToken.getTokenValue())
                .tokenType(OAuth2AccessToken.TokenType.BEARER)
                .expiresIn(expiresAt)
                .scopes(scopes)
                .refreshToken(refreshToken.getTokenValue()).build();
    }

}
